# ==============================================================================================
# BSD 3-Clause Clear License
# Copyright © 2025 ZAMA. All rights reserved.
# ----------------------------------------------------------------------------------------------
# Define a set of utilities command to help with Versal flow
# ==============================================================================================

#################################################
# Common const variables
#################################################
PROJECT_TARGET  := env_var('PROJECT_TARGET')
PROJECT_PART    := env_var('XILINX_PART')
SHELL_VER       := env_var('SHELL_VER')
PROJECT_DIR     := env_var('PROJECT_DIR')

#################################################
# Common path variables
#################################################
WORKDIR         := PROJECT_DIR / "versal"
SCRIPT_DIR      := PROJECT_DIR / "versal/scripts"
OUTDIR          := PROJECT_DIR / "versal/output"

#################################################
# Default values
#################################################
# RTL related
DEFAULT_RTL_FLAGS := "-F TOP_BATCH TOP_BATCH_TOPhpu_BPBS16_TPBS32 -F TOP_PCMAX  TOP_PCMAX_pem2_glwe1_bsk16_ksk16 -F TOP_PC TOP_PC_pem2_glwe1_bsk8_ksk16 -F APPLICATION APPLI_msg2_carry2_pfail64_132b_gaussian_1f72dba -F NTT_MOD NTT_MOD_goldilocks -F NTT_CORE_WMM_ARCH NTT_CORE_ARCH_gf64 -F NTT_CORE_R_PSI NTT_CORE_R2_PSI16 -F NTT_CORE_RDX_CUT NTT_CORE_RDX_CUT_n5c5c1 -F NTT_CORE_DIV NTT_CORE_DIV_1 -F BSK_SLOT_CUT BSK_SLOT8_CUT8 -F KSK_SLOT_CUT KSK_SLOT8_CUT16 -F KSLB KSLB_x2y32z3 -F HPU_PART HPU_PART_delta -F AXI_DATA_W AXI_DATA_W_256 -F FPGA FPGA_v80"

# By default we want to build vivado projects from scratch : from project generation
DEFAULT_STEP := "new"
DEFAULT_STAGE := "all"

DEFAULT_VIVADO_OPT := 'D:MEMORY_FILE_PATH=\"${PROJECT_DIR}/hw/\"'

###################################################################################################
# Recipes
#
# - list_rtl
#     List all the source files via edalize.
#     Today it supports xdc, verilog, systemVerilog and vhdl
#
# - generate_shell
#     Generate a Vivado project including ONLY the block design
#
# - build
#     Build the TOP_IP up to bitstream generation
#     optional arguments can be set to avoid rebuilding everything :
#       * syn:  launches synthesis, if OOM synthesis did not fail we skip it
#       * impl: launches implementation, skips synthesis
#
###################################################################################################

#------------------------------------------------
### Default command
#------------------------------------------------
# List recipes
default:
  just --list

#------------------------------------------------
### List RTL
#------------------------------------------------
# List all rtl modules into a dictionary
list_rtl TOP_IP RTL_FLAGS=DEFAULT_RTL_FLAGS VIVADO_OPT=DEFAULT_VIVADO_OPT : _out-folder
  #!/usr/bin/env bash
  set -euxo pipefail
  python3 {{WORKDIR}}/scripts/declare_define/declare_define.py -l "{{VIVADO_OPT}}" -d {{OUTDIR}}
  # Use run_edalize to generate file list in tcl_dict format
  ${PROJECT_DIR}/hw/scripts/edalize/run_edalize.py -d {{OUTDIR}} -t tcl_dict -m {{TOP_IP}} {{RTL_FLAGS}} --tcl-dict-out {{OUTDIR}}/{{TOP_IP}}_edalize.tcl || exit 1

#------------------------------------------------
### Generate shell
#------------------------------------------------
# Generate only the shell block design along with its sources and verilog wrapper
generate_shell : _out-folder
  #!/usr/bin/env bash
  set -euxo pipefail
  vivado -log "{{OUTDIR}}/create_shell.log" -nojournal -mode batch -source {{SCRIPT_DIR}}/create_shell.tcl

#------------------------------------------------
### Build
#------------------------------------------------
# Print user command
_print_user_cmd TOP_IP STEP RTL_FLAGS VIVADO_OPT:
  #!/usr/bin/env bash
  echo "INFO> Running: just build {{TOP_IP}} {{STEP}} \"{{RTL_FLAGS}}\" \"{{VIVADO_OPT}}\""

# Build Vivado project and launch jobs until bitstream generation
build TOP_IP STEP=DEFAULT_STEP RTL_FLAGS=DEFAULT_RTL_FLAGS VIVADO_OPT=DEFAULT_VIVADO_OPT : (_print_user_cmd TOP_IP STEP RTL_FLAGS VIVADO_OPT) (list_rtl TOP_IP RTL_FLAGS VIVADO_OPT)
  #!/usr/bin/env bash
  set -euxo pipefail
  # these are real exports for inner tcl scripts "build all" and as well "shell" block design generation.
  # we cannot add arguments in nested scripts inside vivado's project mode (i.e below script that is called here).
  export DESIGN_NAME={{TOP_IP}}
  vivado -log {{OUTDIR}}/vivado_integration_{{STEP}}.log -nojournal -mode batch -source {{SCRIPT_DIR}}/build_all.tcl -tclargs {{OUTDIR}} {{STEP}} "{{VIVADO_OPT}}" | tee {{OUTDIR}}/just_build_{{TOP_IP}}.log

program TOP_IP SERIAL_NUMBER STAGE=DEFAULT_STAGE:
  #!/usr/bin/env bash
  TOP_IP={{TOP_IP}}
  STAGE={{STAGE}}
  DRIVERS=("qdma_pf" "ami")

  [[ "$STAGE" != "stage1" && "$STAGE" != "stage2" && "$STAGE" != "all" ]] && echo "[ERROR] please choose a correct stage \"all\", \"stage1\" or \"stage2\"" && exit 1

  # Function to check if a driver is loaded
  check_driver() {
      for driver in "${DRIVERS[@]}"; do
        if lsmod | grep -q "^$driver "; then
            echo "[WARNING] Driver $driver is present, will unload"
            sudo rmmod $driver
        else
            echo "[INFO] Driver $driver already unloaded"
        fi
      done
  }

  if [[ "$STAGE" = "all" || "$STAGE" = "stage1" ]]; then
    echo "[INFO] Stage 1 programming"
    STAGE1_PATH=$(realpath "output/${TOP_IP}_tandem1.pdi")

    # checks
    check_driver
    [ ! -f "$STAGE1_PATH" ] && { echo "[ERROR] $STAGE1_PATH not found" && exit 1; } || { echo "[INFO] Found $STAGE1_PATH"; }

    vivado -mode batch -source scripts/jtag/program_stage1.tcl -tclargs $STAGE1_PATH {{SERIAL_NUMBER}}
  fi

  if [[ "$STAGE" = "all" || "$STAGE" = "stage2" ]]; then
    echo "[INFO] Stage 2 programming"

    PCIE_CARD=$(lspci -d 10ee:50b4)
    DEVICE="${PCIE_CARD%%:*}"
    QMAX="/sys/bus/pci/devices/0000:$DEVICE:00.1/qdma/qmax"
    STAGE2_PATH=$(realpath "output/${TOP_IP}_tandem2.pdi")
    STAGE2_SIZE=$(stat -c%s $STAGE2_PATH)

    # checks
    [ -z "$PCIE_CARD" ] && { echo "[ERROR] V80 PCI device not found" && exit 1; } || { echo "[INFO] Found PCI device"; }
    check_driver
    [ ! -f "$STAGE2_PATH" ] && { echo "[ERROR] $STAGE2_PATH not found" && exit 1; } || { echo "[INFO] Found $STAGE2_PATH"; }
    bootgen -arch versal -read $STAGE2_PATH | grep "rpu_subsystem" && { echo "[INFO] Found elf in stage 2 $STAGE2_PATH"; } || \
    { echo "[ERROR] No elf found in $STAGE2_PATH ... abort" ; exit 1; }

    echo "[INFO] Updating physical functions status"
    sudo bash -c "echo 1 > /sys/bus/pci/devices/0000:$DEVICE:00.0/remove"
    sudo bash -c "echo 1 > /sys/bus/pci/devices/0000:$DEVICE:00.1/remove"
    sudo bash -c "echo 1 > /sys/bus/pci/rescan"
    sleep 2

    echo "[INFO] Initializing queues"
    sudo bash -c "echo 100 > $QMAX"
    dma-ctl qdma${DEVICE}001 q add idx 0 dir h2c
    dma-ctl qdma${DEVICE}001 q start idx 0 dir h2c aperture_sz 4096

    echo "[INFO] Programming stage 2 bitstream"
    dma-to-device -d /dev/qdma${DEVICE}001-MM-0 -f ${STAGE2_PATH} -s ${STAGE2_SIZE} -a 0x000102100000

    echo "[INFO] Reloading drivers and updating pci registers"
    # we must wait a bit for everything to settle down
    sleep 1
    sudo bash -c "echo 1 > /sys/bus/pci/devices/0000:$DEVICE:00.0/remove"
    sudo bash -c "echo 1 > /sys/bus/pci/rescan"
  fi

#------------------------------------------------
### Compile FW
#------------------------------------------------
# Compile ARM fw
compile_fw :
  #!/usr/bin/env bash
  set -euxo pipefail
  # file handling -------------------------------------------------------------
  XSA_PATH={{OUTDIR}}/$SHELL_VER.xsa
  if [ ! -f $XSA_PATH ]; then
    echo "[ERROR]: $XSA_PATH not found"
    exit 1
  fi
  # build ---------------------------------------------------------------------
  cd $PROJECT_DIR/fw/arm/
  ./scripts/build.sh -os freertos10_xilinx -profile v80 -xsa $XSA_PATH
  cp $PROJECT_DIR/fw/arm/build/amc.elf {{OUTDIR}}

#------------------------------------------------
### Merge elf
#------------------------------------------------
# Merge elf from firmware via Vitis and Vivado binaries
merge_elf TOP_IP:
  #!/usr/bin/env bash
  set -euxo pipefail
  # var -----------------------------------------------------------------------
  XSA_PATH={{OUTDIR}}/$SHELL_VER.xsa
  PDI_PATH={{OUTDIR}}/{{TOP_IP}}.pdi
  AMC_PATH={{OUTDIR}}/amc.elf

  TEMPLATE_FILE=$PROJECT_DIR/versal/scripts/fpt/template/pdi_template.json.j2
  OUTPUT_FILE={{OUTDIR}}/pdi_combine.bif

  # file handling -------------------------------------------------------------
  if [ ! -f $XSA_PATH ]; then
    echo "[ERROR]: $XSA_PATH not found"
    exit 1
  fi

  if [ ! -f $PDI_PATH ]; then
    echo "[ERROR]: $PDI_PATH not found"
    exit 1
  fi

  # guard-rail ----------------------------------------------------------------
  # check if an elf is already present
  bootgen -arch versal -read $PDI_PATH | grep "rpu_subsystem" && { echo "[ERROR]: ELF already present in pdi $PDI_PATH"; exit 1; } || echo "[INFO]: no elf found .. continue"

  # save PDI with only FPGA bitstream in a .pdi.noelf file
  cp $PDI_PATH $PDI_PATH.noelf
  # merge elf and top pdi -----------------------------------------------------
  sed -e "s#\^\^top_pdi_path\^\^#$PDI_PATH#g" \
          -e "s#\^\^amc_elf_path\^\^#$AMC_PATH#g" "$TEMPLATE_FILE" > "$OUTPUT_FILE"

  bootgen -arch versal -image $OUTPUT_FILE -w -o $PDI_PATH

#------------------------------------------------
### Merge tandem pdi together
#------------------------------------------------
# this task is meant to be used for two tandem stages
merge_pdi TOP_IP:
  #!/usr/bin/env bash
  set -euxo pipefail
  # var -----------------------------------------------------------------------
  PDI_PATH={{OUTDIR}}/{{TOP_IP}}.pdi
  STAGE_1_PATH={{OUTDIR}}/{{TOP_IP}}_tandem1.pdi
  STAGE_2_PATH={{OUTDIR}}/{{TOP_IP}}_tandem2.pdi

  # file handling -------------------------------------------------------------
  if [[ -f "$STAGE_1_PATH"  && -f "$STAGE_2_PATH" ]]; then
    echo "[INFO] Found stage 1 and 2 files"
  else
    echo "[ERROR] pdi not found"
    exit 1
  fi

  TEMPLATE_FILE=$PROJECT_DIR/versal/scripts/fpt/template/flash_template.json.j2
  OUTPUT_FILE={{OUTDIR}}/pdi_combine.bif

  # merge elf and top pdi -----------------------------------------------------
  sed -e "s#\^\^stage_1_path\^\^#$STAGE_1_PATH#g" \
          -e "s#\^\^stage_2_path\^\^#$STAGE_2_PATH#g" "$TEMPLATE_FILE" > "$OUTPUT_FILE"

  bootgen -arch versal -image $OUTPUT_FILE -w -o $PDI_PATH

#------------------------------------------------
### Rewrite flash
#------------------------------------------------
# Rewrite flash
rewrite_flash:
  #!/usr/bin/env bash
  set -euxo pipefail

  # Checks that the example folder exists, in tcl script, paths are fixed.
  EX_PATH=/opt/amd/aved/amd_v80_gen5x8_23.2_exdes_2_xbtest_stress/flash_setup
  if [ ! -d $EX_PATH ]; then
    echo "[ERROR]: $EX_PATH not found"
    exit 1
  fi

  vivado -log {{OUTDIR}}/rewrite_flash.log -nojournal -mode batch -source {{SCRIPT_DIR}}/jtag/fully_rewrite_flash.tcl

#------------------------------------------------
### Utility commands
#------------------------------------------------
_out-folder:
 #!/usr/bin/env bash
 echo "===================================="
 echo "INFO> output directory :" {{OUTDIR}}
 mkdir -p {{OUTDIR}}
 echo "===================================="
